package com.lwj.framework.security.auth;

import java.util.Collection;

import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.common.exceptions.InvalidTokenException;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;
import org.springframework.security.oauth2.provider.token.store.JwtTokenStore;
import com.lwj.framework.entity.AuthProperties;

public class JwtRedisTokenStore extends JwtTokenStore {
	public final static String STORE_PREFIX = "security_token_";

	private RedisTokenStore redisTokenStore;

	private AuthProperties authProperties;

	public JwtRedisTokenStore(AuthProperties authProperties, JwtAccessTokenConverter jwtTokenEnhancer,
			RedisConnectionFactory redisConnection) {
		super(jwtTokenEnhancer);
		//
		redisTokenStore = new RedisTokenStore(redisConnection);
		redisTokenStore.setPrefix(STORE_PREFIX);
		this.authProperties = authProperties;
	}

	protected InvalidTokenException getInvalidTokenException() {
		return new InvalidTokenException("the token is invalid");
	}

	@Override
	public OAuth2Authentication readAuthentication(OAuth2AccessToken token) {
		OAuth2Authentication jwtAuthentication = super.readAuthentication(token);
		if (jwtAuthentication != null && OAuthUtils.isStore(jwtAuthentication, authProperties)) {// token是否设置了需要存储
			OAuth2Authentication redisAuthentication = redisTokenStore.readAuthentication(token);
			if (redisAuthentication == null) {// token已被销毁
				return null;
			}
		}
		return jwtAuthentication;
	}

	@Override
	public OAuth2Authentication readAuthentication(String token) {
		OAuth2Authentication jwtAuthentication = super.readAuthentication(token);
		if (jwtAuthentication != null && OAuthUtils.isStore(jwtAuthentication, authProperties)) {// token是否设置了需要存储
			OAuth2Authentication redisAuthentication = redisTokenStore.readAuthentication(token);
			if (redisAuthentication == null) {// token已被销毁
				return null;
			}
		}
		return jwtAuthentication;
	}

	@Override
	public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {
		OAuth2AccessToken jwtAccessToken = super.getAccessToken(authentication);
		if (jwtAccessToken != null && OAuthUtils.isStore(authentication, authProperties)) {// token是否设置了需要存储
			OAuth2AccessToken redisAccessToken = redisTokenStore.getAccessToken(authentication);
			if (redisAccessToken == null) {// token已被销毁
				return null;
			}
		}
		return jwtAccessToken;
	}

	@Override
	public OAuth2AccessToken readAccessToken(String tokenValue) {
		OAuth2AccessToken jwtAccessToken = super.readAccessToken(tokenValue);
		if (jwtAccessToken != null && OAuthUtils.isStore(tokenValue)) {// token是否设置了需要存储
			OAuth2AccessToken redisAccessToken = redisTokenStore.readAccessToken(tokenValue);
			if (redisAccessToken == null) {// token已被销毁
				return null;
			}
		}
		return jwtAccessToken;
	}

	@Override
	public OAuth2RefreshToken readRefreshToken(String tokenValue) {
		OAuth2RefreshToken jwtRefreshToken = super.readRefreshToken(tokenValue);
		if (jwtRefreshToken != null && OAuthUtils.isStore(tokenValue)) {// token是否设置了需要存储
			OAuth2RefreshToken redisRefreshToken = redisTokenStore.readRefreshToken(tokenValue);
			if (redisRefreshToken == null) {// token已被销毁
				return null;
			}
		}
		return jwtRefreshToken;
	}

	@Override
	public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication) {
		super.storeAccessToken(token, authentication);
		if (OAuthUtils.isStore(authentication, authProperties)) {
			redisTokenStore.storeAccessToken(token, authentication);
			redisTokenStore.storeRefreshToken(token.getRefreshToken(), authentication);
		}

	}

	@Override
	public void removeAccessToken(OAuth2AccessToken token) {
		super.removeAccessToken(token);
		OAuth2AccessToken redisAccessToken = redisTokenStore.readAccessToken(token.getValue());
		redisTokenStore.removeAccessToken(redisAccessToken);
		redisTokenStore.removeRefreshToken(redisAccessToken.getRefreshToken());
	}

	@Override
	public void storeRefreshToken(OAuth2RefreshToken refreshToken, OAuth2Authentication authentication) {
		super.storeRefreshToken(refreshToken, authentication);
		if (OAuthUtils.isStore(authentication, authProperties)) {
			redisTokenStore.storeRefreshToken(refreshToken, authentication);
		}
	}

	@Override
	public void removeRefreshToken(OAuth2RefreshToken token) {
		super.removeRefreshToken(token);
		redisTokenStore.removeRefreshToken(token);
	}

	@Override
	public Collection<OAuth2AccessToken> findTokensByClientId(String clientId) {
		Collection<OAuth2AccessToken> jwtAccessTokens = super.findTokensByClientId(clientId);
		if (jwtAccessTokens != null && !jwtAccessTokens.isEmpty()
				&& OAuthUtils.isStoreClient(clientId, authProperties)) {// token是否设置了需要存储
			Collection<OAuth2AccessToken> redisAccessTokens = redisTokenStore.findTokensByClientId(clientId);
			if (redisAccessTokens == null) {// token已被销毁
				return null;
			}
		}
		return jwtAccessTokens;
	}

	@Override
	public Collection<OAuth2AccessToken> findTokensByClientIdAndUserName(String clientId, String userName) {
		Collection<OAuth2AccessToken> jwtAccessTokens = super.findTokensByClientIdAndUserName(clientId, userName);
		if (jwtAccessTokens != null && !jwtAccessTokens.isEmpty()
				&& OAuthUtils.isStoreClient(clientId, authProperties)) {// token是否设置了需要存储
			Collection<OAuth2AccessToken> redisAccessTokens = redisTokenStore.findTokensByClientIdAndUserName(clientId,
					userName);
			if (redisAccessTokens == null) {// token已被销毁
				return null;
			}
		}
		return jwtAccessTokens;
	}

	@Override
	public OAuth2Authentication readAuthenticationForRefreshToken(OAuth2RefreshToken token) {
		OAuth2Authentication jwtAuthentication = super.readAuthenticationForRefreshToken(token);
		if (jwtAuthentication != null && OAuthUtils.isStore(jwtAuthentication, authProperties)) {// token是否设置了需要存储
			OAuth2Authentication redisAuthentication = redisTokenStore.readAuthenticationForRefreshToken(token);
			if (redisAuthentication == null) {// token已被销毁
				return null;
			}
		}
		return jwtAuthentication;
	}

	@Override
	public void removeAccessTokenUsingRefreshToken(OAuth2RefreshToken refreshToken) {
		super.removeAccessTokenUsingRefreshToken(refreshToken);
		redisTokenStore.removeAccessTokenUsingRefreshToken(refreshToken);
	}
}
